#!/usr/bin/env python
# coding: utf-8

# In[ ]:


from __future__ import absolute_import, division, print_function

import numpy as np
import tensorflow as tf
import pandas as pd
import random
import math




import tensorflow.keras as K

from tensorflow_probability import distributions as tfd

from tensorflow.keras.layers import Input, Dense, Activation, Concatenate
from tensorflow.keras.callbacks import EarlyStopping, TensorBoard, ReduceLROnPlateau

from sklearn.linear_model import LinearRegression
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import MinMaxScaler

import matplotlib as mpl

import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
import seaborn as sns
import statsmodels.api as sm
from scipy import stats
from sklearn.metrics import r2_score as r_2
from scipy.stats import ranksums

from scipy.stats import wilcoxon
from scipy.stats import ks_2samp
from scipy.stats import mannwhitneyu
from scipy.stats import combine_pvalues
from csv import writer

mpl.rcParams.update(mpl.rcParamsDefault)
mpl.rcParams.update({'font.size': 12})


import warnings
warnings.filterwarnings("always")

def remove_ax_window(ax):
    """
        Remove all axes and tick params in pyplot.
        Input: ax object.
    """
    ax.spines["top"].set_visible(False)    
    ax.spines["bottom"].set_visible(False)    
    ax.spines["right"].set_visible(False)    
    ax.spines["left"].set_visible(False)  
    ax.tick_params(axis=u'both', which=u'both',length=0)
    
dpi = 140
x_size = 8
y_size = 4
alt_font_size = 14

save_figure = False
use_tb = False


# In[ ]:


class MDN(tf.keras.Model):

    def __init__(self, neurons=100, components = 2):
        super(MDN, self).__init__(name="MDN")
        self.neurons = neurons
        self.components = components
        
        self.h1 = Dense(neurons, activation="relu", name="h1")
        self.h2 = Dense(neurons, activation="relu", name="h2")
        
        self.alphas = Dense(components, activation="softmax", name="alphas")
        self.mus = Dense(components, name="mus")
        self.sigmas = Dense(components, activation="nnelu", name="sigmas")
        self.pvec = Concatenate(name="pvec")
        
    def call(self, inputs):
        x = self.h1(inputs)
        x = self.h2(x)
        
        alpha_v = self.alphas(x)
        mu_v = self.mus(x)
        sigma_v = self.sigmas(x)
        
        return self.pvec([alpha_v, mu_v, sigma_v])
    

    

def nnelu(input):
    """ Computes the Non-Negative Exponential Linear Unit
    """
    return tf.add(tf.constant(1, dtype=tf.float32), tf.nn.elu(input))

def slice_parameter_vectors(parameter_vector):
    """ Returns an unpacked list of paramter vectors.
    """
    return [parameter_vector[:,i*components:(i+1)*components] for i in range(no_parameters)]

def gnll_loss(y, parameter_vector):
    """ Computes the mean negative log-likelihood loss of y given the mixture parameters.
    """
    alpha, mu, sigma = slice_parameter_vectors(parameter_vector) # Unpack parameter vectors
    
    gm = tfd.MixtureSameFamily(
        mixture_distribution=tfd.Categorical(probs=alpha),
        components_distribution=tfd.Normal(
            loc=mu,       
            scale=sigma))
    
    log_likelihood = gm.log_prob(tf.transpose(y)) # Evaluate log-probability of y
    
    return -tf.reduce_mean(log_likelihood, axis=-1)

tf.keras.utils.get_custom_objects().update({'nnelu': Activation(nnelu)})


# # Travel time prediction

# In[ ]:


filename="training"
#Create a numpy array containing items in csv file
import csv

with open(filename+'.csv', newline='') as csvfile:
    csvreader = csv.reader(csvfile, delimiter=' ', quotechar='|')
    
    i=0
    for row in csvreader:
        if i==0:
            i = i +1
        else:
            L = row[0].split(',')
            
            carId = int(L[0])
            currTime = float(L[1])
            roadStartPt = int(L[2])
            roadEndPt = int(L[3])
            roadLength = float(L[4])
            roadGradient = float(L[5])
            carDesiredSpeed = float(L[6])
            vehType = int(L[7])
            carSpeed = float(L[8])
            carAcceleration = float(L[9])
            carLength = float(L[10])
            avgCarLength = float(L[11])
            numOfCarsOnRoad = int(L[12])
            vehicleInputRate = float(L[13])
            maxNumOfCarsOnEdge = int(L[14])
            numOfLanesOnRoad=int(L[15])
            travelTime = float(L[16])
            
            
            
            if i==1:
                ARR=np.array([[carId, currTime, roadStartPt, roadEndPt, roadLength, roadGradient, carDesiredSpeed, vehType,
                                carSpeed, carAcceleration, carLength, avgCarLength, numOfCarsOnRoad, 
                                vehicleInputRate,maxNumOfCarsOnEdge, numOfLanesOnRoad,travelTime]])
                i = i+1
            else:
                ARR=np.append(ARR,np.array([[carId, currTime, roadStartPt, roadEndPt, roadLength, roadGradient, carDesiredSpeed, vehType, 
                                                carSpeed, carAcceleration, carLength, avgCarLength, numOfCarsOnRoad, 
                                                vehicleInputRate,maxNumOfCarsOnEdge, numOfLanesOnRoad,travelTime]]),axis=0)
            
            del carId, currTime, roadStartPt, roadEndPt, roadLength, roadGradient, carDesiredSpeed 
            del vehType,carSpeed, carAcceleration, carLength, avgCarLength, numOfCarsOnRoad 
            del vehicleInputRate,maxNumOfCarsOnEdge, numOfLanesOnRoad,travelTime

print(ARR)  

import pickle
with open(filename+'_2023'+'.data', 'wb') as filehandle:
    pickle.dump(ARR, filehandle)
del ARR


# In[ ]:





# In[ ]:



#Set parameters for model
model_opt ='adam'
model_no_parameters = 3
model_components = 10 #number of distributions in this model
model_neurons = 30
numOfEpochs=10000

opt = model_opt
no_parameters = model_no_parameters
components = model_components #number of distributions in this model
neurons = model_neurons


import pickle
with open(filename+'_2023'+'.data', 'rb') as filehandle:
    ARR = pickle.load(filehandle)
    
    



#split data for each road separately
ARR_allEdges = {}
edgesInNetwork = []

for item in ARR:
    edge_start = item[2]
    edge_end = item[3]
    
    if (edge_start,edge_end) in edgesInNetwork:
        ARR_oneEdge = ARR_allEdges[(edge_start,edge_end)]
        ARR_oneEdge.append(item)
        
    else:
        
        ARR_oneEdge = []
        ARR_oneEdge.append(item)
        ARR_allEdges[(edge_start,edge_end)]=ARR_oneEdge
        edgesInNetwork.append((edge_start,edge_end))


#convert the arrays to numpy array
ARR_allEdges_numpy = {}
for (edge_start,edge_end) in edgesInNetwork:
    ARR_oneEdge = ARR_allEdges[(edge_start,edge_end)]
    ARR_oneEdge2 = np.asarray(ARR_oneEdge).astype(object)
    ARR_allEdges_numpy[(edge_start,edge_end)] = ARR_oneEdge2
    

#train each mdn for each edge #In a loop
mdn_dict = {}
mdn_maxMin_dict = {}
dnn_dict = {}
min_max_scaler_x_dict={}
min_max_scaler_y_dict={}

for (edge_start,edge_end) in edgesInNetwork:
    
    ARR = ARR_allEdges_numpy[(edge_start,edge_end)]


    

    arrayMean = np.mean(ARR, axis=0)
    avgCarLength_pick = arrayMean[10]
    avgCarLength_pick = np.round(avgCarLength_pick,3)
    avgCarLength_col = [avgCarLength_pick]*ARR.shape[0]

    X1 = ARR[:,[4,5,6,7]]  
    X2=np.asarray(avgCarLength_col).astype(np.float64)
    X3 = ARR[:,[12,13,14,15]]
    X=np.column_stack((X1,X2,X3))

    y = ARR[:,[16]]


    y=np.asarray(y).astype(np.float64)
    
    ARR = np.column_stack((X,y))
    ARR_allEdges_numpy[(edge_start,edge_end)] = ARR

    
    x_data = X.copy()
    y_data = y.copy()

    min_max_scaler_x = MinMaxScaler()
    min_max_scaler_y = MinMaxScaler()
    y_data = min_max_scaler_y.fit_transform(y_data)
    x_data = min_max_scaler_x.fit_transform(x_data)
    
    
    
        
    x_train = x_data.copy()
    y_train = y_data.copy()

    
    
    y_min = float(np.amin(y_train))
    y_max = float(np.amax(y_train))


  
    

    mdn_2 = MDN()
    mdn_2 = MDN(neurons=neurons, components=components)
    mdn_2.compile(loss=gnll_loss, optimizer=opt)

    

    mon = EarlyStopping(monitor='val_loss', min_delta=0, patience=5, verbose=0, mode='auto')

    
    if len(x_data)>10:
        mdn_2.fit(x=x_train, y=y_train,epochs=numOfEpochs, callbacks=[mon], batch_size=128, verbose=0,validation_split=0.1)
        
    else:
        mdn_2.fit(x=x_train, y=y_train,epochs=numOfEpochs, callbacks=[mon], batch_size=128, verbose=0)

    mdn_dict[(edge_start,edge_end)] = mdn_2
    mdn_maxMin_dict[(edge_start,edge_end)] = (y_min, y_max)
    min_max_scaler_x_dict[(edge_start,edge_end)]= min_max_scaler_x
    min_max_scaler_y_dict[(edge_start,edge_end)]= min_max_scaler_y
    
    
    del ARR,X,y,x_data,y_data,min_max_scaler_x,min_max_scaler_y,x_train, y_train,y_min,y_max,mdn_2,mon



print("train each mdn for each edge #In a loop - complete")   

   


# In[ ]:


#TEST



#filename = 'testing'


# In[ ]:



import heapq

import csv
timefilename = '_TimeData'
carlengthfilename= '_averageCarLengthData'
carpropertiesfilename= '_CarGeneralProperties'
edgepropertiesfilename= '_edgePropertiesData'
pathsfilename = '_paths'



Paths={}
PathIdList=[]
pathProbabilityList=[]
with open(pathsfilename+'.csv', newline='') as csvfile:
    csvreader = csv.reader(csvfile, delimiter=' ', quotechar='|')
    
    i=0
    for row in csvreader:
        if i==0:
            i = i +1
        else:
            L = row[0].split(',')
            
            
            pathId = int(L[0])
            pathprob = float(L[1])
            PathIdList.append(pathId)
            pathProbabilityList.append(pathprob)
            pathdat = []
            
            for i in range(2,len(L),1):
                #print("L[i]=",L[i])
                if L[i]!='':
                    dcpt = int(L[i])
                    pathdat.append(dcpt)
            
            
            Paths[pathId]=pathdat




carDesSpeed = 50
vehType = 100

with open(carpropertiesfilename+'.csv', newline='') as csvfile:
    csvreader = csv.reader(csvfile, delimiter=' ', quotechar='|')
    
    i=0
    for row in csvreader:
        if i==0:
            i = i +1
        else:
            L = row[0].split(',')
            
            carDesSpeed = float(L[0])
            carDesSpeed = (carDesSpeed*5)/18
            vehType = int(L[1])
            
            
#read edge properties
edgeProperties={}
numOfCarsOnRoad = {}
roadQofCars={}
waitQofCars={}
with open(edgepropertiesfilename+'.csv', newline='') as csvfile:
    csvreader = csv.reader(csvfile, delimiter=' ', quotechar='|')
    
    i=0
    for row in csvreader:
        if i==0:
            i = i +1
        else:
            L = row[0].split(',')
            
            roadStartPt = int(L[0])
            roadEndPt = int(L[1])
            roadLength = float(L[2])
            roadGradient = float(L[3])
            maxNumOfCarsOnEdge = int(L[4])
            numOfLanesOnThisRoad = int(L[5])
            edgeProperties[(roadStartPt,roadEndPt)]=[roadLength,roadGradient,maxNumOfCarsOnEdge,numOfLanesOnThisRoad]
            numOfCarsOnRoad[(roadStartPt,roadEndPt)]=0
            roadQofCars[(roadStartPt,roadEndPt)]=[]
            waitQofCars[(roadStartPt,roadEndPt)]=[]
            
            if (roadStartPt,None) not in roadQofCars:
                roadQofCars[(roadStartPt,None)]=[]
            if (roadEndPt,None) not in roadQofCars:
                roadQofCars[(roadEndPt,None)]=[]
            
            
# read other car length data
avgCarLength_read = 0
vehInputRate = 0
totNumOfCars=0
totcarlengths=0

with open(carlengthfilename+'.csv', newline='') as csvfile:
    csvreader = csv.reader(csvfile, delimiter=' ', quotechar='|')
    
    i=0
    for row in csvreader:
        if i==0:
            i = i +1
        elif i==1:
            L = row[0].split(',')
            avgCarLength_read = float(L[0])
            totNumOfCars = int(L[1])
            totcarlengths = float(L[2])
            vehInputRate = float(L[3])
            

#time data 
simulationTime=0
vissimRunTime=0
with open(timefilename+'.csv', newline='') as csvfile:
    csvreader = csv.reader(csvfile, delimiter=' ', quotechar='|')
    
    i=0
    for row in csvreader:
        if i==0:
            i = i +1
        elif i==1:
            L = row[0].split(',')
            simulationTime = float(L[0])
            vissimRunTime = float(L[1])


# In[ ]:







with open('totTrvlmdn'+'.csv', 'w',newline='') as f_object:
    writer_object = writer(f_object)
    writer_object.writerow(["MDN total travel time"])  
    f_object.close()





# number of cars in 1 hour = vehInputRate
NumOfCars = int((vehInputRate/(60*60))*simulationTime)
meanInterArrivalTime = 3600/vehInputRate

carid = 0
carStartTimes=[]
car_sortTime = []



while carid < NumOfCars:
    random_num_exp = np.random.exponential(meanInterArrivalTime, int(vehInputRate))# (mean,num of samples)
    carid = carid + 1
    prev=0
    if len(carStartTimes)>0:
        prev=carStartTimes[(len(carStartTimes)-1)][1]
    carStartTimes.append((carid,round((prev+ random_num_exp[0]),2)))
    
   
    i=1
    while (carid < NumOfCars) and (i<len(random_num_exp)):
        carid = carid + 1
        prev=carStartTimes[(len(carStartTimes)-1)][1]
        carStartTimes.append((carid,round((prev+ random_num_exp[i]),2)))
        i = i+ 1

    

for carStartTimes_id in range(0,len(carStartTimes)):
    pathid = random.choices(PathIdList, pathProbabilityList)[0]
    
    heapq.heappush(car_sortTime, (carStartTimes[carStartTimes_id][1], [carStartTimes[carStartTimes_id][0],
                                                                       carStartTimes[carStartTimes_id][1],
                                                                       carStartTimes[carStartTimes_id][1],
                                                                       pathid,0]))
    




    
#run simulation of cars in priority queue
while len(car_sortTime) >0:
    
    [carid,startTime,arrivalTime,pathid,carPathPos] = heapq.heappop(car_sortTime)[1]
    currNode = Paths[pathid][carPathPos]
    prevNode = None
    nextNode = None
    nextnextNode = None
    
    if len(Paths[pathid])>int(carPathPos)+1:
        nextNode = Paths[pathid][int(carPathPos)+1]
    if len(Paths[pathid])>int(carPathPos)+2:
        nextnextNode = Paths[pathid][int(carPathPos)+2]
    if carPathPos>0:
        prevNode = Paths[pathid][int(carPathPos)-1]
        
    if nextNode==None:
        
        if carPathPos>0:
            numOfCarsOnRoad[(prevNode,currNode)] = numOfCarsOnRoad[(prevNode,currNode)] - 1
        
        
        totTrvlTime= arrivalTime - carStartTimes[carid-1][1]
        with open('totTrvlmdn'+'.csv', 'a',newline='') as f_object:
            writer_object = writer(f_object)
            writer_object.writerow([totTrvlTime])  
            f_object.close()
        
        continue
    
    
        
    
        
    
    
    maxLimitCarNums = edgeProperties[(currNode,nextNode)][2]
    
    if numOfCarsOnRoad[(currNode,nextNode)] < maxLimitCarNums:
        
        #update cars on exiting road 
        if carPathPos>0:
            
            numOfCarsOnRoad[(prevNode,currNode)] = numOfCarsOnRoad[(prevNode,currNode)] - 1
            
            
        
        if carPathPos<(len(Paths[pathid])-1): 
            inputForModel = [edgeProperties[(currNode,nextNode)][0],edgeProperties[(currNode,nextNode)][1],
                             carDesSpeed,vehType,avgCarLength_read,numOfCarsOnRoad[(currNode,nextNode)],
                             vehInputRate,edgeProperties[(currNode,nextNode)][2],edgeProperties[(currNode,nextNode)][3]]
            inputForModel= [inputForModel]

            min_max_scaler_x = min_max_scaler_x_dict[(currNode,nextNode)]
            min_max_scaler_y = min_max_scaler_y_dict[(currNode,nextNode)]

            inputForModel = min_max_scaler_x.transform(inputForModel)

            inputForModel = np.array(inputForModel)
            mdn_2 = mdn_dict[(currNode,nextNode)]
            (y_min, y_max) = mdn_maxMin_dict[(currNode,nextNode)]
            alpha, mu, sigma = slice_parameter_vectors(mdn_2.predict(inputForModel))
            gm = tfd.MixtureSameFamily(
                    mixture_distribution=tfd.Categorical(probs=alpha),
                    components_distribution=tfd.Normal(
                        loc=mu,       
                        scale=sigma))


            pred_travelTime = gm.sample()

            while (pred_travelTime<y_min) or (pred_travelTime>y_max):
                pred_travelTime = gm.sample()

            pred_travelTime = min_max_scaler_y.inverse_transform([pred_travelTime.numpy()])

            pred_travelTime = pred_travelTime[0][0]
            pred_travelTime = round(pred_travelTime,2)
            pred_arrivalTime = startTime+pred_travelTime
            pred_arrivalTime = round(pred_arrivalTime,2)

            
            

            if not (nextnextNode==None):
                maxLimitCarNums_nextedge = edgeProperties[(nextNode,nextnextNode)][2]
                if numOfCarsOnRoad[(nextNode,nextnextNode)] < maxLimitCarNums_nextedge: 
                    heapq.heappush(car_sortTime, (pred_arrivalTime, [carid,pred_arrivalTime,pred_arrivalTime,pathid,carPathPos+1]))
                    
                else:
                    print("no space in edge (",nextNode,",",nextnextNode,") for car ",carid)
                    break
                    
            else:
                heapq.heappush(car_sortTime, (pred_arrivalTime, [carid,pred_arrivalTime,pred_arrivalTime,pathid,carPathPos+1]))
                
            

            numOfCarsOnRoad[(currNode,nextNode)] = numOfCarsOnRoad[(currNode,nextNode)] + 1
            
            
        else:
            print(carid, " reached destination")
            
            
            
            


            
    
    
    
    
    del carid,carPathPos,currNode
    


            


        


# In[ ]:




